package app

import (
	"chatgpt/ai"
	"chatgpt/ai/chatglm"
	openai2 "chatgpt/ai/openai"
	"chatgpt/controller/father"
	"chatgpt/helper"
	"chatgpt/helper/code"
	"chatgpt/logq"
	"chatgpt/services/completion"
	"github.com/gin-gonic/gin"
	"github.com/slclub/glog"
	"github.com/slclub/go-tips"
	"strings"
	"time"
)

func QuickHandle(ctx *gin.Context) {
	choice := ctx.Query("model")
	switch choice {
	case "completion", "", "completions":
		CompletionHandle(ctx)
	case "chat":
		ChatHandle(ctx)
	}
}

func CompletionHandle(ctx *gin.Context) {
	question := father.RawString(ctx, "q")   //ctx.Query("q")
	mod := father.RawString(ctx, "model")    //ctx.Query("model")
	unafe := father.RawString(ctx, "unsafe") //ctx.Query("unsafe")
	if !completion.CheckQuestion(question) {
		helper.Echo(ctx, int(code.DATA_NIL), "")
		return
	}
	if !ai.AppAiEneity.CheckQuestion(question) && !ai.AppAiEneity.Unsafe(question, unafe) {
		helper.Echo(ctx, int(code.SUCCESS), "抱歉没有找到与“"+question+"”相关的信息")
		return
	}
	uid := (helper.Uid(ctx))
	queue := completion.Mgr().Get(uid)
	// 缓存 查到一直问相同的问题，从而直接回答
	if queue != nil {
		answer := queue.RepeatCompletion(question)
		if answer != "" {
			helper.Echo(ctx, int(code.SUCCESS), answer)
			return
		}
	}
	stream := father.RawAny(ctx, "stream")
	stream_str := ""
	switch val := stream.(type) {
	case bool:
		if val == true {
			stream_str = "yes"
		}
	case string:
		if val != "false" && val != "" {
			stream_str = "true"
		}
	}
	smodel := father.RawString(ctx, "smodel")
	if smodel == "" {
		smodel = "chatglm_6b"
	}
	mcode, replay := ai.Completion(ctx, question, map[string]string{
		"model":  mod,
		"smodel": smodel,
		"stops":  completion.Mgr().StopSequence(uid),
		"unsafe": helper.UNSAFE_MOD,
		"stream": stream_str,
	}, func() string {
		nquestion, _ := chatglm.Search.CustomQuestion(question)
		if queue != nil {
			//nquestion = queue.ConvPrompt(question)
		}
		return nquestion
	})
	if replay == nil {
		//helper.Echo(ctx, int(code.DATA_GPT_ERROR), nil)
		//return
		glog.Infof("App.Completion.GPT.error %v:%v", code.DATA_GPT_ERROR, code.CODENO(code.DATA_GPT_ERROR).String())
		replay = &helper.ResponseDesc{MsgCode: mcode}
	}

	if mcode == code.SUCCESS && replay != nil {
		completion.Mgr().Add(uid, completion.NewDialogue(question, tips.String(replay.Data)))
	}
	if stream_str == "" && replay != nil {
		helper.Echo(ctx, mcode, completion.ErrorCodeCompletion(tips.String(replay.Data), mcode), replay.Tokens)
	}

	gptRecord := &openai2.AppGPT{}
	writeRes := gptRecord.SetRecordLog(map[string]any{
		"timestamp": time.Now().Unix(),
		"userAgent": ctx.GetHeader("User-Agent"),
		"uid":       uid,
		"smodel":    ctx.Query("smodel"),
		"question":  question,
		"answer":    replay.Data.(string),
		"ip":        ctx.RemoteIP(),
		"date":      time.Now(),
	})

	if writeRes != true {
		logq.DebugF("HANDLE.COMPLETION uid:%v ip:%v tokens:%v prompt-all:%s", uid, ctx.RemoteIP(), replay.Tokens, writeRes)
	}

	logq.InfoF("HANDLE.COMPLETION code:%v uid:%v ip:%v tokens:%v prompt:%s", int(mcode), uid, ctx.RemoteIP(), replay.Tokens, question)
}

func ChatHandle(ctx *gin.Context) {
	question := father.RawString(ctx, "q")        //ctx.Query("q")
	mod := father.RawString(ctx, "model")         //ctx.Query("model")
	smodel := father.RawString(ctx, "smodel")     //ctx.Query("smodel")
	unafe := father.RawString(ctx, "unsafe")      //ctx.Query("unsafe")
	stream_str := father.RawString(ctx, "stream") //ctx.Query("unsafe")
	if !completion.CheckQuestion(question) {
		helper.Echo(ctx, int(code.DATA_NIL), "")
		return
	}
	if !openai2.AppChatGPT.CheckQuestion(question) && !openai2.AppChatGPT.Unsafe(question, unafe) {
		helper.Echo(ctx, int(code.SUCCESS), "抱歉没有找到与“"+question+"”相关的信息")
		return
	}
	uid := (helper.Uid(ctx))
	queue := completion.Mgr().Get(uid)
	nquestion := question
	if queue != nil {
		nquestion = queue.ConvPromptToChat(question)
	} else {
		completion.Mgr().EmptyQueue.Uid = uid
		nquestion = completion.Mgr().EmptyQueue.ConvPromptToChat(question)
	}

	if strings.Trim(smodel, " ") == "" {
		smodel = helper.Config().Sub("ChatGPT").Sub("chat").GetString("model")
	}
	mcode, replay := ai.Chat(ctx, nquestion, map[string]any{
		"mod":    mod,
		"smodel": smodel,
		"stream": stream_str,
		"uid":    uid,
	}, nil)

	if mcode == code.SUCCESS {
		completion.Mgr().Add(uid, completion.NewDialogue(question, tips.String(replay.Data)))
	}

}

func CompletionClearHandle(ctx *gin.Context) {
	uid := (helper.Uid(ctx))
	queue := completion.Mgr().Get(uid)
	if queue == nil {
		helper.Echo(ctx, int(code.SUCCESS), "")
		return
	}

	queue.Clear()
	helper.Echo(ctx, int(code.SUCCESS), "服务会话清理")
}
